# -*- coding: utf-8 -*- 
# @Time : 2022/4/4 16:33 
# @Author : zzuxyj 
# @File : 16-model-test.py
import torch
import torchvision
from PIL import Image

"""
测试模型0
"""
# 数据加载
img_path = "../dataset/imgs/horse.png"
img = Image.open(img_path)
# png格式图片要进行转换
img = img.convert("RGB")
print(img)
# img.show()

# 转换格式
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32,32)),
    torchvision.transforms.ToTensor()
])
img = transform(img)
#print(img.shape)

#GPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 加载训练好的模型
model = torchvision.models.vgg16(pretrained=False)
model.classifier[6] = torch.nn.Linear(4096, 10)
modelParam = torch.load("./model/vgg16-model19-CIFAR10.pth" , map_location=device)
model.load_state_dict(modelParam)
model.to(device)

#print(model)
# 图片输入格式改变
img = torch.reshape(img , (1,3,32,32))
img = img.to(device)
# target
target = ["airplane" , "automobile" , "bird" , "cat" , "deer" , "dog" , "frog" , "horse" , "ship" ,"truck" ]
target_ZH_CN = ["飞机" ," 汽车","鸟","猫","鹿","狗","蛙","马","船","卡车"]
#测试
model.eval()
with torch.no_grad():
    output = model(img)
    print(target_ZH_CN[output.argmax(1)])


